import numpy

import wlc

def wlc_force(kb_t_over_p, k_lc, z_over_lc):
    kb_t_over_p_over_k_lc = kb_t_over_p / k_lc
    c = (
        (4 * kb_t_over_p_over_k_lc + 4) / k_lc ** 2,
        -(12 * kb_t_over_p_over_k_lc + 8) / k_lc * z_over_lc + (9 * kb_t_over_p_over_k_lc + 8) / k_lc,
        (12 * kb_t_over_p_over_k_lc + 4) * z_over_lc ** 2 - (18 * kb_t_over_p_over_k_lc + 8) * z_over_lc + 6 * kb_t_over_p_over_k_lc + 4,
        -4 * kb_t_over_p * z_over_lc ** 3 + 9 * kb_t_over_p * z_over_lc ** 2 -  6 * kb_t_over_p * z_over_lc
    )
    roots = numpy.roots(c)
    # TODO: Do we really need to enforce it to be non-negative?
    y_threshold = max(0, k_lc * (z_over_lc - 1))
    valid_roots = list(frozenset([r.real for r in roots if r.imag == 0 and r.real >= y_threshold]))
    assert len(valid_roots) == 1, valid_roots
    r = valid_roots[0]
    assert r <= k_lc * z_over_lc
    return r

def discretise_wlc(wlc_ascent, time_step):
    kb_t_over_p = float(wlc_ascent.kb_t_over_p)
    k_lc = float(wlc_ascent.k_lc)
    v_over_lc = float(wlc_ascent.v_over_lc)
    min_f = float(wlc_ascent.min_f)
    max_f = float(wlc_ascent.max_f)
    min_t = wlc.f_inverse(kb_t_over_p, k_lc, v_over_lc, min_f)
    max_t = wlc.f_inverse(kb_t_over_p, k_lc, v_over_lc, max_f)
    time_points = list(numpy.arange(min_t, max_t, time_step))
    # For consistency with linear and wlc models.
    if time_points[-1] < max_t:
        time_points.append(max_t)
    result = list((t, wlc_force(kb_t_over_p, k_lc, v_over_lc * t)) for t in time_points)
    # TODO: assert result[0][1] about the same as min_f and similar for last element.
    return result

